import numpy as np

from core.variable import Variable
from core.cuda import cuda
from utils.functions_collect import matmul
import matplotlib.pyplot as plt


# Generate toy dataset
np.random.seed(0)
x = np.random.rand(100, 1)
y = 5 + 2 * x + np.random.rand(100, 1)
x, y = Variable(x), Variable(y)

W = Variable(np.zeros((1, 1)))
b = Variable(np.zeros(1))


def predict(x):
    y = matmul(x, W) + b
    return y


def mean_squared_error(x0, x1):
    diff = x0 - x1
    return sum(diff ** 2) / len(diff)


lr = 0.1
iters = 100

for i in range(iters):
    y_pred = predict(x)
    loss = mean_squared_error(y, y_pred)

    W.cleargrad()
    b.cleargrad()
    loss.backward()

    W.data = cuda.to_array(W.data)
    W.grad.data = cuda.to_array(W.grad.data)
    b.data = cuda.to_array(b.data)
    b.grad.data = cuda.to_array(b.grad.data)

    # Update .data attribute (No need grads when updating params)
    W.data -= lr * W.grad.data
    b.data -= lr * b.grad.data
    print(W, b, loss)


# Plot
x.data = cuda.to_numpy(x.data)
y.data = cuda.to_numpy(y.data)

plt.scatter(x.data, y.data, s=10)
plt.xlabel('x')
plt.ylabel('y')
y_pred = predict(x)
y_pred.data = cuda.to_numpy(y_pred.data)
plt.plot(x.data, y_pred.data, color='r')
plt.show()